from openpyxl import load_workbook
from algorithm.propagation_model.IC_model import ICModel
from algorithm.seed_nodes.PRTH import PRTH
from algorithm.seed_nodes.community_PR import community_PR


def diploma_COPRTH(G, communities, k=1):
    H = G.copy()
    seed_nodes = PRTH(H)
    # 加载工作簿
    wb = load_workbook('output/diploma.xlsx')
    # 选择要操作的工作表
    ws = wb['COBase']
    for i in range(1, 6):
        seed_size = i * k
        seed_set = community_PR(communities, seed_size, seed_nodes)
        # 使用选取的种子节点进行信息传播并输出影响力
        model = ICModel(G)
        influence = model.simulate(seed_set, 0.1)
        data = len(influence)
        print(f"COPRTH算法的种子节点个数为{seed_size},该组种子节点的影响力为{data}")
        # 要添加数据的单元格
        cell = ws.cell(row=i + 2, column=4)
        # 添加数据
        cell.value = data
    # 保存更改
    wb.save('output/diploma.xlsx')